import ctypes
import numpy as np
import logging

class RGA:
    def __init__(self,lib_path,width,height):
        """
        lib_path::
        """
        lib = ctypes.cdll.LoadLibrary(lib_path)
        # prototypes
        lib.nv12_to_rgb_init.argtypes = [ctypes.c_int, ctypes.c_int]
        lib.nv12_to_rgb_init.restype = ctypes.c_int

        lib.nv12_alloc_src.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_void_p)]
        lib.nv12_alloc_src.restype = ctypes.c_int  # returns allocated size

        lib.nv12_free_src.argtypes = [ctypes.c_int, ctypes.c_size_t, ctypes.c_void_p]
        lib.nv12_free_src.restype = None

        lib.nv12_to_rgb_convert.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int]
        lib.nv12_to_rgb_convert.restype = ctypes.c_int

        lib.nv12_get_dst_fd.restype = ctypes.c_int
        lib.nv12_get_dst_virt.restype = ctypes.c_void_p
        lib.nv12_get_dst_size.restype = ctypes.c_size_t

        lib.nv12_to_rgb_deinit.restype = None
        self.lib=lib
        self.WIDTH=width
        self.HEIGHT=height
        self.init_dma()

    def init_dma(self):
        # 初始化（例如 3840x2160 支持 4K）
        if self.lib.nv12_to_rgb_init(self.WIDTH, self.HEIGHT) != 0:
            raise RuntimeError("init failed")
        # 分配 src NV12
        w, h = self.WIDTH, self.HEIGHT
        self.src_fd = ctypes.c_int()
        self.src_virt = ctypes.c_void_p()
        self.allocated_size = self.lib.nv12_alloc_src(w, h, ctypes.byref(self.src_fd), ctypes.byref(self.src_virt))
        if self.allocated_size <= 0:
            raise RuntimeError("alloc src failed")
        logging.info("src_fd:", self.src_fd.value, "src_virt:", hex(self.src_virt.value), "size:", self.allocated_size)
    
    def convert_nv12_to_rgb(self,data_buffer):
        buf_type = ctypes.c_char * self.allocated_size
        buf = buf_type.from_address(self.src_virt.value)
        buf[:len(data_buffer)] = data_buffer

        # 调用转换
        ret = self.lib.nv12_to_rgb_convert(self.src_fd.value, self.WIDTH, self.HEIGHT)
        if ret != 0:
            raise RuntimeError("convert failed")
        # 读取 dst
        dst_fd = self.lib.nv12_get_dst_fd()
        dst_ptr = self.lib.nv12_get_dst_virt()
        dst_size = self.lib.nv12_get_dst_size()
        # 将 RGB 数据保存 (注意：RGB 不含头，仅 raw rgb24)
        dst_buf_type = ctypes.c_ubyte * dst_size
        dst_buf = dst_buf_type.from_address(dst_ptr)
        # 生成 np.ndarray，注意RGB24，每像素3字节
        # rgb_array = np.ctypeslib.as_array(dst_buf)
        # rgb_array = rgb_array.reshape((self.HEIGHT, self.WIDTH, 3))
        return dst_buf
    
    def release(self):
        # 清理
        self.lib.nv12_free_src(self.src_fd.value, self.allocated_size, self.src_virt.value)
        self.lib.nv12_to_rgb_deinit()